# -*- coding: utf-8 -*-
"""
观察输入数据效果并发布到ROS进行RVIZ可视化
输入：rgb图、深度图、实例分割mask图
输出：
    - rgb图像
    - 实例分割mask叠加在rgb上
    - 深度图转换的点云（原始颜色渲染）
"""

import rospy
import cv2
import numpy as np
import open3d as o3d
from PIL import Image as PilImage
from sensor_msgs.msg import Image, PointCloud2, PointField
from cv_bridge import CvBridge
from std_msgs.msg import Header

import tkinter as tk
import time

from common import get_color


class DataVisualizer:
    def __init__(self):
        rospy.init_node("data_visualizer", anonymous=True)

        # 创建图像桥接器
        self.bridge = CvBridge()

        # 发布者
        self.rgb_pub = rospy.Publisher("/output/rgb_image", Image, queue_size=10)
        self.overlay_pub = rospy.Publisher("/output/overlay_mask", Image, queue_size=10)
        self.pointcloud_pub = rospy.Publisher("/output/point_cloud", PointCloud2, queue_size=10)
        self.instance_pointcloud_pub = rospy.Publisher("/output/instance_point_cloud", PointCloud2, queue_size=10)

        # 内参假设（需根据相机参数调整）
        self.fx = 384.281 # 焦距x
        self.fy = 383.770  # 焦距y
        self.cx = 329.505  # 中心x
        self.cy = 242.564  # 中心y

        self.file_index = 0
        self.start_gui()

    def start_gui(self):
        """启动 GUI 线程"""
        import threading
        gui_thread = threading.Thread(target=self.create_gui)
        gui_thread.daemon = True
        gui_thread.start()

    def create_gui(self):
        """创建 GUI 界面，支持上下箭头调整 file_index"""
        root = tk.Tk()
        root.title("File Index 控制")

        label = tk.Label(root, text="输入 file_index:")
        label.pack()

        entry_var = tk.StringVar(value=str(self.file_index))  # 与输入框绑定
        entry = tk.Entry(root, textvariable=entry_var)
        entry.pack()

        def update_file_index_from_entry():
            try:
                new_index = int(entry_var.get())
                self.file_index = new_index
                rospy.loginfo(f"已更新 file_index 到: {new_index}")
            except ValueError:
                rospy.logwarn("请输入有效的整数！")
                entry_var.set(str(self.file_index))  # 恢复原值

        def on_increase():
            self.file_index += 1
            entry_var.set(str(self.file_index))
            rospy.loginfo(f"已更新 file_index 到: {self.file_index}")

        def on_decrease():
            self.file_index -= 1
            entry_var.set(str(self.file_index))
            rospy.loginfo(f"已更新 file_index 到: {self.file_index}")

        button_frame = tk.Frame(root)
        button_frame.pack(pady=5)

        down_button = tk.Button(button_frame, text="⬇", width=2, command=on_decrease)
        down_button.pack(side=tk.LEFT, padx=2)

        up_button = tk.Button(button_frame, text="⬆", width=2, command=on_increase)
        up_button.pack(side=tk.LEFT, padx=2)

        confirm_button = tk.Button(root, text="更新", command=update_file_index_from_entry)
        confirm_button.pack(pady=5)

        root.mainloop()

    def publish_image(self, image, topic_name):
        msg = self.bridge.cv2_to_imgmsg(image, encoding="bgr8")
        msg.header.stamp = rospy.Time.now()
        getattr(self, f"{topic_name}_pub").publish(msg)

    def overlay_mask_on_rgb(self, rgb_image, combined_mask, alpha=0.3):
        """将mask叠加到RGB图像上"""
        mask_image = np.zeros_like(rgb_image)
        unique_instance_ids = np.unique(combined_mask)

        # 用于保存所有有效 mask 的布尔掩码（排除 id == 0）
        valid_mask = np.zeros_like(combined_mask, dtype=bool)

        # 为每个类别生成不同颜色的mask，并记录有效区域
        has_mask = False
        for id in unique_instance_ids:
            if id == 0:
                continue  # 跳过背景
            
            has_mask = True
            color = get_color(id)
            color = (color[2], color[1], color[0])
            current_mask = (combined_mask == id)
            mask_image[current_mask] = color
            valid_mask |= current_mask  # 合并所有有效的 mask

        # 仅对有效 mask 区域进行叠加
        overlayed_image = rgb_image.copy()
        if has_mask:
            overlayed_image[valid_mask] = cv2.addWeighted(rgb_image[valid_mask], 1, mask_image[valid_mask], alpha, 0)
        return overlayed_image

    def depth_to_pointcloud(self, rgb_image, depth_image, depth_scale=1.0):
        """
        将深度图转为点云，使用RGB图像着色
        如果提供mask_image，则只对mask区域生成点云
        """
        h, w = depth_image.shape
        i, j = np.indices((h, w))

        # 计算相机坐标系下的xyz坐标
        z = depth_image * depth_scale
        x = (j - self.cx) * z / self.fx
        y = (i - self.cy) * z / self.fy

        points = np.stack((x, y, z), axis=-1)

        # 使用RGB颜色
        rgb_image = cv2.cvtColor(rgb_image, cv2.COLOR_BGR2RGB)
        colors = rgb_image.reshape(-1, 3)
        points = points.reshape(-1, 3)

        return self.convert_to_PointCloud2(points, colors)

    def depth_to_mask_color_pointcloud(self, depth_image, mask_image, depth_scale=1.0):
        """
        将深度图转为点云，使用mask的像素值作为颜色
        """
        h, w = depth_image.shape
        i, j = np.indices((h, w))

        # 计算相机坐标系下的xyz坐标
        z = depth_image * depth_scale
        x = (j - self.cx) * z / self.fx
        y = (i - self.cy) * z / self.fy

        points = np.stack((x, y, z), axis=-1)

        # 提取mask值
        mask_values = mask_image.flatten()
        unique_ids = np.unique(mask_values[mask_values > 0])  # 跳过背景（0）

        # 为每个实例分配固定颜色
        color_map = {}
        for idx in unique_ids:
            color = get_color(idx)
            color_map[idx] = [color[0], color[1], color[2]]

        colors = np.array([color_map.get(val, [0, 0, 0]) for val in mask_values])

        points = points.reshape(-1, 3)

        return self.convert_to_PointCloud2(points, colors)
    
    def convert_to_PointCloud2(self, points, colors):
        """
        将 numpy 数组转换为 sensor_msgs.PointCloud2 消息。
        使用 NumPy 向量化操作加速。
        """
        start_time = time.time()
        header = Header()
        header.stamp = rospy.Time.now()
        header.frame_id = "world"

        # 定义 PointCloud2 字段
        fields = [
            PointField(name='x', offset=0, datatype=PointField.FLOAT32, count=1),
            PointField(name='y', offset=4, datatype=PointField.FLOAT32, count=1),
            PointField(name='z', offset=8, datatype=PointField.FLOAT32, count=1),
            PointField(name='rgb', offset=16, datatype=PointField.UINT32, count=1),
        ]

        num_points = len(points)
        point_step = 32
        row_step = point_step * num_points

        # 构建结构化 dtype（注意：必须按 4 字节对齐）
        dtype = np.dtype([
            ('x', np.float32),
            ('y', np.float32),
            ('z', np.float32),
            ('pad1', np.uint8, 4),  # 占位 4 字节 padding
            ('rgb', np.uint32),
            ('pad2', np.uint8, 12),  # 占位 12 字节 padding
        ])

        # 创建结构化数组并填充数据
        structured_array = np.zeros(num_points, dtype=dtype)
        structured_array['x'] = points[:, 0]
        structured_array['y'] = points[:, 1]
        structured_array['z'] = points[:, 2]

        # 转换颜色格式 (r, g, b) -> uint32
        r = colors[:, 0].astype(np.uint32)
        g = colors[:, 1].astype(np.uint32)
        b = colors[:, 2].astype(np.uint32)
        a = np.full_like(r, 255)
        structured_array['rgb'] = (a << 24) | (r << 16) | (g << 8) | b

        # 转为字节流
        data = structured_array.tobytes()

        result = PointCloud2(
            header=header,
            height=1,
            width=num_points,
            fields=fields,
            is_bigendian=False,
            point_step=point_step,
            row_step=row_step,
            data=data,
            is_dense=True
        )

        end_time = time.time()
        print(f"convert_to_PointCloud2 time: {end_time - start_time} s")
        return result

    def run(self):
        rate = rospy.Rate(1)

        # 示例路径（请替换为实际路径）
        file_dir = "/mnt/data/3d-lidar/semantic/self_collect/for_panoptic_mapping"

        while not rospy.is_shutdown():
            file_index_str = f"{self.file_index:06d}"
            rgb_path = f"{file_dir}/{file_index_str}_color.png"
            depth_path = f"{file_dir}/{file_index_str}_depth.tiff"
            mask_path = f"{file_dir}/{file_index_str}_predicted.png"
            # 读取图像
            rgb_image = cv2.imread(rgb_path)
            mask_image = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
            depth_image = np.array(PilImage.open(depth_path))

            # 发布原始RGB图像
            self.publish_image(rgb_image, "rgb")

            # 生成mask叠加图像并发布
            overlay_image = self.overlay_mask_on_rgb(rgb_image, mask_image, alpha=1.0)
            self.publish_image(overlay_image, "overlay")

            # 发布RGB着色点云
            pointcloud_rgb_msg = self.depth_to_pointcloud(rgb_image, depth_image)
            self.pointcloud_pub.publish(pointcloud_rgb_msg)

            # 发布Mask着色点云
            pointcloud_mask_msg = self.depth_to_mask_color_pointcloud(depth_image, mask_image)
            self.instance_pointcloud_pub.publish(pointcloud_mask_msg)

            rate.sleep()


if __name__ == "__main__":
    visualizer = DataVisualizer()
    visualizer.run()